# -*- coding: utf-8 -*-
"""
文件名: monitor.py

功能概述:
    - 监控训练输出目录中的 checkpoint-<step> 子目录（HF/LlamaFactory 训练风格）。
    - 逐步对这些 checkpoint 做推理评估（调用 scripts/vllm_infer.py 子进程）。
    - 评估结果按 step 输出到 eval_out_dir/eval_step_<step>.jsonl。
    - 我们按“安全 checkpoint”顺序慢慢吃：
        * 排序所有 checkpoint（升序）
        * 最新的那个当作“还在写的，别碰”
        * 其他的都是“安全候选”
        * 在这些安全候选里，从最小 step 开始，找还没评过的，依次评
        * 每轮最多评 max_eval_per_round 个（默认 1）
    - 评测完成后，保留最近 retain_k 个 checkpoint（step 最大的 retain_k 个）；更老的会被删掉，释放磁盘。
      ⚠ 注意：只有在评测成功时才会触发删除。评测失败则不会删，方便你调试。

关键异常策略 (这是这版的重点改动):
    1. 如果一次评测失败（推理子进程异常 / 超时 / 没生成结果文件等）:
        - 不删除任何 checkpoint
        - 不把该 step 记入 evaluated_steps
        - 立即将错误信息写入 eval_out_dir/eval_step_<step>.error.txt
        - 立刻退出整个监控进程 (sys.exit(1))
          -> 你的训练不会被我们杀，但监控停下来提醒你“出事了”。
    2. 如果一次评测成功:
        - 记录到 state.evaluated_steps
        - 清理老的 checkpoint，最多保留 retain_k 个最近的
        - 继续下一轮循环

健康退出逻辑:
    - run_infer_with_healthcheck() 会同步等待子进程，直到:
        * 子进程自然退出
        * 或我们检测到 save_name.jsonl 已经稳定写好 >= quiet_secs 秒，于是我们主动 terminate/kill 子进程
        * 或 timeout_sec 触发，强杀并视为失败
    - run_infer_with_healthcheck() 返回 (ok: bool, err_msg: Optional[str])
        * ok=True: 视为评测成功
        * ok=False: err_msg 中有详细调试信息

使用样例:
    python fundus_reasoner/eval/monitor.py \
        --watch_dir ./experiments/ckpt/sft_stage1_simpleQA \
        --eval_out_dir ./experiments/eval_output \
        --dataset_dir /path/to/Alpaca_data_monitor \
        --dataset_name monitor_2000 \
        --vllm_infer_py scripts/vllm_infer.py \
        --poll_interval 10 \
        --timeout_sec 300 \
        --quiet_secs 10 \
        --retain_k 3 \
        --top_p 0.9 \
        --top_k 1 \
        --python_bin /path/to/python

注意事项 / 约定:
    - watch_dir 下面既可能有若干 checkpoint-<step> 子目录，也可能有根目录自己的一份 "最终权重" (model-00001-of-000xx.safetensors 等)。
      当前逻辑仍然只会自动评测 checkpoint-* 这些子目录；根目录最终权重不在自动评估范围内。
      你可以训练结束后手动跑一次评估这个最终权重。
    - CUDA 资源/显卡占用:
        我们这里不会并发多个评估。run_infer_with_healthcheck 是同步阻塞的。
        如果你希望评估永远跑在固定单卡，比如 GPU7，可以在下面的 subprocess.Popen(...) 那里补 env={"CUDA_VISIBLE_DEVICES": "7"}。
        目前默认继承当前进程环境。
    - state 文件:
        eval_out_dir/_watch_state.json 用来记录哪些 step 已经评测过。
        如果评测失败，我们不会把这个 step 写进去，这样你修完 bug 后重新起 monitor，还是会从这个 step 继续尝试。

实现细节:
    - 我们新增了小异常类 EvalRunError，用于在 main() 里区分“评测失败”这种我们希望立刻退出的致命情况。
"""


import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import time
from dataclasses import dataclass, asdict
from typing import List, Dict, Optional, Tuple
import shlex
import glob

# ===================== 自定义异常 =====================

class EvalRunError(Exception):
    """表示一次 checkpoint 评测失败，需要立即终止整个监控进程。"""
    pass


# ===================== 参数解析（集中默认值） =====================

def parse_args():
    parser = argparse.ArgumentParser(description="监控训练Checkpoint并触发vLLM评测")

    # 训练权重目录，包含 checkpoint-* 子目录
    parser.add_argument(
        "--watch_dir",
        "-i",
        type=str,
        default="./experiments/ckpt/sft_stage1_simpleQA",
        help="训练保存checkpoint的根目录（包含 checkpoint-* 子目录）"
    )

    # 评测输出目录
    parser.add_argument(
        "--eval_out_dir",
        "-o",
        type=str,
        default="./experiments/eval_output/sub_exp_1",
        help="评测结果输出目录（保存 eval_step_<N>.jsonl 以及错误日志）"
    )

    # 推理所需测试数据集
    parser.add_argument(
        "--dataset_dir",
        type=str,
        default="/home/zhangpinglu/data0/gy/Dataset/fundusreasoner/Alpaca_data_monitor",
        help="评测数据集根目录（alpaca_dir）"
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
        default="monitor_2000",
        help="数据集名称（例如 monitor_2000 或 med_pub_test_compare）"
    )

    # 推理脚本和解释器
    parser.add_argument(
        "--vllm_infer_py",
        type=str,
        default="scripts/vllm_infer.py",
        help="vLLM 推理脚本路径"
    )
    parser.add_argument(
        "--python_bin",
        type=str,
        default=sys.executable,
        help="Python 解释器路径，用于启动子进程 (会继承当前环境的CUDA_VISIBLE_DEVICES 等)"
    )

    # 轮询与健康监测
    parser.add_argument(
        "--poll_interval",
        type=int,
        default=10,
        help="主循环监控轮询间隔秒"
    )
    parser.add_argument(
        "--timeout_sec",
        type=int,
        default=6000,
        help="单次评测子进程的全局超时时间（秒）"
    )
    parser.add_argument(
        "--quiet_secs",
        type=int,
        default=60,
        help="结果文件静默时间（秒），在这段时间里文件未再增长 -> 认为写完，可主动终止推理子进程"
    )
    parser.add_argument(
        "--stable_secs",
        type=int,
        default=30,
        help="checkpoint 目录需要静默多长时间才视为‘安全可评测’（秒）"
    )

    # 保留策略
    parser.add_argument(
        "--retain_k",
        type=int,
        default=3,
        help="最多保留最近K个checkpoint目录（按 step 从大到小选K个保留）。只有在评测成功后才触发清理"
    )

    # 生成/采样参数（传给 vllm_infer.py）
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
        help="采样温度 (一般要求0)"
    )
    parser.add_argument(
        "--top_p",
        type=float,
        default=0.9,
        help="top_p"
    )
    parser.add_argument(
        "--top_k",
        type=int,
        default=1,
        help="top_k"
    )

    # 限速：一轮最多评多少个安全 checkpoint
    parser.add_argument(
        "--max_eval_per_round",
        type=int,
        default=1,
        help="每次主循环中，最多评测多少个 checkpoint（串行逐个评测，默认1）"
    )

    return parser.parse_args()


# ===================== 内部数据结构 =====================

CKPT_PAT = re.compile(r"^checkpoint-(\d+)$")


@dataclass
class WatchState:
    """
    监控状态:
      - evaluated_steps: 已成功评测的 step 列表
      - last_seen_latest_step: 上一轮看到的“最新 step”
    """
    evaluated_steps: List[int]
    last_seen_latest_step: Optional[int] = None

    @staticmethod
    def load(path: str) -> "WatchState":
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                d = json.load(f)
            return WatchState(**d)
        return WatchState(evaluated_steps=[])

    def save(self, path: str) -> None:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(asdict(self), f, ensure_ascii=False, indent=2)


# ===================== 基础工具函数 =====================

def list_checkpoints(watch_dir: str) -> List[Tuple[int, str, float]]:
    """
    扫描 watch_dir 下的 checkpoint-* 子目录，返回:
        [(step, abs_path, mtime), ...]  按 step 升序排列
    说明:
        - step 从目录名解析出来 (checkpoint-<step>)
        - mtime 使用该目录本身的最近修改时间
    """
    if not os.path.isdir(watch_dir):
        # raise NotADirectoryError(f"watch_dir 不是目录: {watch_dir}")
        os.mkdir(watch_dir, exist_ok=True)
    out: List[Tuple[int, str, float]] = []
    for name in os.listdir(watch_dir):
        m = CKPT_PAT.match(name)
        if not m:
            continue
        step = int(m.group(1))
        p = os.path.join(watch_dir, name)
        if os.path.isdir(p):
            try:
                mtime = os.path.getmtime(p)
            except Exception as e:
                raise RuntimeError(f"获取 mtime 失败: {p} | err={e}")
            out.append((step, p, mtime))

    out.sort(key=lambda x: x[0])  # step 升序
    return out


def is_stable(path: str, stable_secs: int) -> bool:
    """
    判断某个 checkpoint 目录在最近 stable_secs 秒内有没有发生修改。
    只有稳定的 checkpoint 我们才会尝试评测。
    """
    try:
        mtime = os.path.getmtime(path)
    except Exception as e:
        raise RuntimeError(f"获取 mtime 失败: {path} | err={e}")

    return (time.time() - mtime) >= stable_secs


def build_save_name(eval_out_dir: str, step: int) -> str:
    """
    构造当前 step 的结果文件名。统一放在 eval_out_dir 下。
    """
    os.makedirs(eval_out_dir, exist_ok=True)
    return os.path.join(eval_out_dir, f"eval_step_{step}.jsonl")


def build_error_log_name(eval_out_dir: str, step: int) -> str:
    """
    如果评测失败，记录错误日志的位置。
    """
    os.makedirs(eval_out_dir, exist_ok=True)
    return os.path.join(eval_out_dir, f"eval_step_{step}.error.txt")


# ===================== 推理子进程封装 =====================


def run_infer_with_healthcheck(
    python_bin: str,
    vllm_infer_py: str,
    model_path: str,
    dataset_dir: str,
    dataset_name: str,
    save_name: str,
    timeout_sec: int,
    quiet_secs: int,
    temperature: float,
    top_p: float,
    top_k: int,
) -> Tuple[bool, str]:
    """
    启动 vLLM 推理子进程，并实时监控输出文件 save_name 的生成情况。
    同时捕获 stdout/stderr 到文件，以便调试。
    """
    cmd = [
        python_bin, vllm_infer_py,
        "--model_name_or_path", model_path,
        "--dataset_dir", dataset_dir,
        "--dataset", dataset_name,
        "--save_name", save_name,
        "--template", "qwen2_vl",
        "--temperature", str(temperature),
        "--top_p", str(top_p),
        "--top_k", str(top_k),
    ]
    cmd_str = " ".join(shlex.quote(x) for x in cmd)

    # 派生日志文件路径
    base_log = f"{save_name}.log"
    stdout_log = f"{base_log}.stdout"
    stderr_log = f"{base_log}.stderr"

    try:
        with open(stdout_log, "w", encoding="utf-8") as fout, \
             open(stderr_log, "w", encoding="utf-8") as ferr:
            proc = subprocess.Popen(
                cmd,
                stdout=fout,
                stderr=ferr,
                text=True,
            )
    except Exception as e:
        return False, (
            f"启动推理子进程失败: {' '.join(cmd)} | "
            f"err={type(e).__name__}: {e}"
        )

    start_ts = time.time()
    last_seen_mtime = None
    last_change_ts = None

    # 实时捕获输出流（写文件时同步刷新）
    outs_buf = []
    errs_buf = []

    while True:
        time.sleep(1.0)

        # 每秒刷新 stdout/stderr 内容
        if proc.stdout:
            line = proc.stdout.readline()
            if line:
                outs_buf.append(line)
                with open(stdout_log, "a", encoding="utf-8") as f:
                    f.write(line)
        if proc.stderr:
            line = proc.stderr.readline()
            if line:
                errs_buf.append(line)
                with open(stderr_log, "a", encoding="utf-8") as f:
                    f.write(line)

        # 1️⃣ 超时检查
        elapsed = time.time() - start_ts
        if elapsed > timeout_sec:
            try:
                proc.terminate()
                time.sleep(2)
                if proc.poll() is None:
                    proc.kill()
            except Exception:
                pass
            return False, (
                f"推理子进程超时(>{timeout_sec}s)，模型={model_path}, save_name={save_name}"
            )

        # 2️⃣ 输出文件检查
        if os.path.exists(save_name):
            try:
                mtime = os.path.getmtime(save_name)
            except Exception as e:
                try:
                    proc.terminate()
                    time.sleep(2)
                    if proc.poll() is None:
                        proc.kill()
                except Exception:
                    pass
                return False, (
                    f"结果文件存在但无法读取 mtime: {save_name} | err={e}"
                )

            # 文件变化检测
            if (last_seen_mtime is None) or (mtime != last_seen_mtime):
                last_seen_mtime = mtime
                last_change_ts = time.time()
            else:
                if last_change_ts is not None and (time.time() - last_change_ts) >= quiet_secs:
                    # 文件静默时间达到阈值 -> 尝试优雅退出
                    proc.terminate()
                    try:
                        out_tail, err_tail = proc.communicate(timeout=5)
                    except subprocess.TimeoutExpired:
                        proc.kill()
                        out_tail, err_tail = proc.communicate()

                    # 最后补全日志
                    with open(stdout_log, "a", encoding="utf-8") as f:
                        f.write(out_tail or "")
                    with open(stderr_log, "a", encoding="utf-8") as f:
                        f.write(err_tail or "")

                    return True, ""

        # 3️⃣ 子进程自然退出
        rc = proc.poll()
        if rc is not None:
            out_tail, err_tail = proc.communicate()

            # 保存最后输出
            with open(stdout_log, "a", encoding="utf-8") as f:
                f.write(out_tail or "")
            with open(stderr_log, "a", encoding="utf-8") as f:
                f.write(err_tail or "")

            # ✅ 成功A: 正常退出 + 文件存在
            if rc == 0 and os.path.exists(save_name):
                return True, ""

            # ✅ 成功B: returncode != 0 但结果文件存在
            if rc != 0 and os.path.exists(save_name):
                sys.stderr.write(
                    "[warn] 子进程非零退出码，但结果文件已生成，视为成功。\n"
                    f"  returncode={rc}\n"
                    f"  save_name={save_name}\n"
                    f"  stderr_tail:\n{(err_tail or '')[-2000:]}\n"
                )
                return True, ""

            # ❌ 失败
            return False, (
                "推理子进程异常退出，且未生成结果文件.\n"
                f"cmd={' '.join(cmd)}\n"
                f"returncode={rc}\n"
                f"stderr_tail:\n{(err_tail or '')[-2000:]}\n"
                f"stdout_tail:\n{(out_tail or '')[-2000:]}"
            )



# ===================== checkpoint 清理策略 =====================
def prune_checkpoints_incremental_safe(
    ckpts: List[Tuple[int, str, float]],
    retain_k: int,
    evaluated_steps: List[int],
) -> None:
    """
    仅删除“已成功评测过”的旧 checkpoint。
    保留最近 retain_k 个 step 最大的（不论是否评测）。
    对未评测的 checkpoint 一律不删除。
    """
    if retain_k <= 0:
        return
    if len(ckpts) <= retain_k:
        return

    # 按 step 升序
    ckpts_sorted = sorted(ckpts, key=lambda x: x[0])

    # 找出所有已评测的 checkpoint
    evaluated_ckpts = [c for c in ckpts_sorted if c[0] in evaluated_steps]

    if len(evaluated_ckpts) <= retain_k:
        return

    # 要删除的 = 多余的旧评测结果
    num_to_delete = len(evaluated_ckpts) - retain_k
    to_delete = evaluated_ckpts[:num_to_delete]

    for step, path, _ in to_delete:
        try:
            shutil.rmtree(path)
            sys.stdout.write(f"[info] 已删除评测过的旧 checkpoint: step={step} path={path}\n")
        except Exception as e:
            sys.stderr.write(
                f"[warn] 删除 checkpoint 失败: step={step} path={path} | err={e}\n"
            )

# ===================== 主循环 =====================
def main():
    """
    监控主循环（已修改）：

    行为总结：
      - 如果 state.can_exit 为 True：将监控目录下的所有 checkpoint 视为安全候选（包括最新的那个），
        串行从最小 step 开始评测，直到把监控目录下所有 checkpoint 都评测完毕，然后调用 ring.py 并返回。
      - 如果 state.can_exit 为 False（默认）：保持原有行为：
        * 当 ckpts < 2 时不评测（没有“安全的上一份”）
        * 否则把最后一个当作 still-writing，不评它；只对 safe_candidates 中 mtime 静默 >= stable_secs 的项评测
    """
    args = parse_args()
    os.makedirs(args.eval_out_dir, exist_ok=True)
    state_path = os.path.join(args.eval_out_dir, "_watch_state.json")
    state = WatchState.load(state_path)

    sys.stdout.write(
        f"[boot] watch_dir={args.watch_dir}\n"
        f" eval_out_dir={args.eval_out_dir}\n"
        f" dataset={args.dataset_name} from {args.dataset_dir}\n"
        f" poll_interval={args.poll_interval}s, timeout={args.timeout_sec}s, quiet={args.quiet_secs}s\n"
        f" retain_k={args.retain_k}\n"
    )

    while True:
        try:
            # ---------- 0) 检查结束标志文件 ----------
            try:
                safetensor_found = bool(glob.glob(os.path.join(args.watch_dir, "*.safetensor")))
                config_found = (
                    os.path.exists(os.path.join(args.watch_dir, "config.json"))
                    or os.path.exists(os.path.join(args.watch_dir, "configs.json"))
                )
            except Exception as e:
                raise RuntimeError(f"检测结束标志文件失败: {args.watch_dir} | err={e}")

            if safetensor_found or config_found:
                if not getattr(state, "can_exit", False):
                    state.can_exit = True
                    state.save(state_path)
                    sys.stdout.write(
                        f"[info] 监控目录检测到结束标志: safetensor={safetensor_found} config={config_found} -> state.can_exit=True\n"
                    )

            # ---------- 1) 列出 checkpoints ----------
            ckpts = list_checkpoints(args.watch_dir)  # 返回按 step 升序的 [(step, path, mtime), ...]
            if not ckpts:
                # 没有任何 checkpoint，等待
                time.sleep(args.poll_interval)
                continue

            # ---------- 2) 决定哪些是“安全候选” ----------
            if getattr(state, "can_exit", False):
                # 仅在 can_exit 时：把所有 ckpts 视为安全（包括最新一个），
                # 并且不再依赖 is_stable（因为训练已完成，我们要尽快把所有评完）
                safe_candidates = [(s, p) for s, p, _ in ckpts]
            else:
                # 常规模式：如果 checkpoint 数量 < 2，则没有“安全的上一份”可以评
                if len(ckpts) < 2:
                    time.sleep(args.poll_interval)
                    continue
                # 把最后一个当作 still-writing，不评它
                safe_candidates = [(s, p) for s, p, _ in ckpts[:-1]]

            # ---------- 3) 过滤稳定（非 can_exit 模式下需要 stable） ----------
            stable_candidates: List[Tuple[int, str]] = []
            if getattr(state, "can_exit", False):
                # can_exit 模式下不检查 mtime，直接全部视为 stable
                stable_candidates = list(safe_candidates)
            else:
                for step, path in safe_candidates:
                    if is_stable(path, args.stable_secs):
                        stable_candidates.append((step, path))

            # ---------- 4) 选择本轮要评测的 batch ----------
            to_eval_batch: List[Tuple[int, str]] = []
            for step, path in stable_candidates:
                if step not in state.evaluated_steps:
                    to_eval_batch.append((step, path))
                if len(to_eval_batch) >= args.max_eval_per_round:
                    break

            # ---------- 5) 如果本轮没有新任务，检查是否可以退出（当 can_exit 且 所有 ckpts 都已经评测完） ----------
            if not to_eval_batch:
                latest_step = ckpts[-1][0]
                if state.last_seen_latest_step != latest_step:
                    state.last_seen_latest_step = latest_step
                    state.save(state_path)

                if getattr(state, "can_exit", False):
                    all_steps = [s for s, _, _ in ckpts]
                    evaluated_set = set(state.evaluated_steps)
                    if all(s in evaluated_set for s in all_steps):
                        msg = f"监控目录 {args.watch_dir} 检测到结束标志且所有 checkpoint 已评测完毕"
                        sys.stdout.write(f"[finish] {msg}\n")
                        try:
                            subprocess.run([args.python_bin, "ring.py", "-t", msg], check=True)
                        except Exception as e:
                            raise RuntimeError(f"调用 ring.py 失败: err={e}")
                        return

                time.sleep(args.poll_interval)
                continue

            # ---------- 6) 串行执行 to_eval_batch ----------
            for step, ckpt_path in to_eval_batch:
                save_name = build_save_name(args.eval_out_dir, step)
                error_log = build_error_log_name(args.eval_out_dir, step)

                if os.path.exists(save_name):
                    sys.stdout.write(f"[skip] 结果已存在，跳过评测: step={step} -> {save_name}\n")
                    if step not in state.evaluated_steps:
                        state.evaluated_steps.append(step)
                        state.save(state_path)
                    continue

                sys.stdout.write(f"[run ] 评测 checkpoint: step={step} path={ckpt_path}\n")

                ok, err_msg = run_infer_with_healthcheck(
                    python_bin=args.python_bin,
                    vllm_infer_py=args.vllm_infer_py,
                    model_path=ckpt_path,
                    dataset_dir=args.dataset_dir,
                    dataset_name=args.dataset_name,
                    save_name=save_name,
                    timeout_sec=args.timeout_sec,
                    quiet_secs=args.quiet_secs,
                    temperature=args.temperature,
                    top_p=args.top_p,
                    top_k=args.top_k,
                )

                if not ok:
                    # 写错误日志（失败不记录为 evaluated），并中断监控
                    try:
                        with open(error_log, "w", encoding="utf-8") as f:
                            f.write(
                                "[EVAL FAILED]\n"
                                f"step={step}\n"
                                f"ckpt={ckpt_path}\n"
                                "err_msg=\n"
                                f"{err_msg}\n"
                            )
                    except Exception as e_w:
                        raise RuntimeError(f"写入错误日志失败: {error_log} | err={e_w}")

                    sys.stderr.write(
                        f"[fail] 评测失败: step={step}\n"
                        f" ckpt={ckpt_path}\n"
                        f" err={err_msg[:1000]}\n"
                        f" 已写入错误日志: {error_log}\n"
                        " 监控程序将退出，请手动检查问题后重启。\n"
                    )
                    raise EvalRunError(f"评测失败(step={step}), 监控退出; err简述: {err_msg[:200]}")

                # 成功：记录状态并触发 prune（prune 只在评测成功时执行）
                sys.stdout.write(f"[done] 评测完成: step={step} -> {save_name}\n")
                if step not in state.evaluated_steps:
                    state.evaluated_steps.append(step)
                    state.save(state_path)

                prune_checkpoints_incremental_safe(
                    ckpts=ckpts,
                    retain_k=args.retain_k,
                    evaluated_steps=state.evaluated_steps
                )

            # ---------- 7) 本轮完成后，记录最新 step 并再次检查是否可退出 ----------
            latest_step = ckpts[-1][0]
            if state.last_seen_latest_step != latest_step:
                state.last_seen_latest_step = latest_step
                state.save(state_path)

            if getattr(state, "can_exit", False):
                all_steps = [s for s, _, _ in ckpts]
                evaluated_set = set(state.evaluated_steps)
                if all(s in evaluated_set for s in all_steps):
                    msg = f"监控目录 {args.watch_dir} 检测到结束标志且所有 checkpoint 已评测完毕"
                    sys.stdout.write(f"[finish] {msg}\n")
                    try:
                        subprocess.run([args.python_bin, "ring.py", "-t", msg], check=True)
                    except Exception as e:
                        raise RuntimeError(f"调用 ring.py 失败: err={e}")
                    return

        except EvalRunError as ee:
            sys.stderr.write(f"[fatal] {ee}\n")
            sys.exit(1)
        except Exception as e:
            sys.stderr.write(f"[fatal] [主循环异常] {type(e).__name__}: {e}\n")
            sys.exit(1)

        time.sleep(args.poll_interval)


# ===================== 入口 =====================

if __name__ == "__main__":
    main()
